
import dota_utils2 as util
import os
import numpy as np
from PIL import Image

def yolo2dota(imgpath, txtpath, dstpath, extractclassname):
    """
    :param imgpath: the path of images
    :param txtpath: the path of txt in YOLO format
    :param dstpath: the path of txt in dota format
    :param extractclassname: the category you selected
    :return:
    """
    filelist = util.GetFileFromThisRootDir(txtpath)
    for fullname in filelist:
        print(fullname)

        objects = util.parse_dota_hbb(fullname)

        name = os.path.splitext(os.path.basename(fullname))[0]
        img_fullname = os.path.join(imgpath, name + '.png')
        img = Image.open(img_fullname)
        img_w, img_h = img.size
        # print img_w,img_h
        with open(os.path.join(dstpath, name + '.txt'), 'w') as f_out:
            for obj in objects:
                hbb = obj['hbb']
                if len(hbb) != 4:
                    continue
                poly = util.hbb2poly(hbb, img_w, img_h)
                if obj['class'] < len(extractclassname):
                    name = extractclassname[obj['class']]
                else:
                    continue
                outline = ' '.join(list(map(str, poly))) + ' ' + name + ' ' + '0'
                f_out.write(outline + '\n')


if __name__ == '__main__':
    ## an example
    yolo2dota(r'..\datasets\dota\images',
                 r'..\datasets\dota\labels',
                 r'..\datasets\dota\test',
              util.wordname_15)

